from typing import List
from math import sin, cos, log2
from numpy import fft, pi
import matplotlib.pyplot as plt

xs:List[float] = [0.3535, 0.3535, 0.6464, 1.0607, 0.3535, -1.0607, -1.3535, -0.3535]

def Euler(theta:float)->complex:
	real:float = cos(theta)
	imag:float = sin(theta)
	return complex(real, imag)

def RecursiveFFT(src:List[float])->List[complex]:
	n = len(src)
	if n == 1:
		return [complex(a) for a in src]
	wn = Euler(-2 * pi / n)
	w = complex(1)
	evenList = src[::2]
	oddList = src[1::2]
	evenRet = RecursiveFFT(evenList)
	oddRet = RecursiveFFT(oddList)
	ret = [None] * n
	for k in range(n//2):
		t = w * oddRet[k]
		ret[k] = evenRet[k] + t
		ret[k+n//2] = evenRet[k] - t
		w = w * wn
	return ret

def ReverseBit(num:int, bitCount:int)->int:
	binary = bin(num)
	reverse = binary[-1:1:-1]
	reverse = reverse + '0' * (bitCount - len(reverse))
	return int(reverse, 2)

def IterativeFFT(src:List[float])->List[complex]:
	n = len(src)
	bitCount = int(log2(n))
	rev:List[float] = [src[ReverseBit(i, bitCount)] for i in range(n)]
	# 对每一层
	for s in range(1, bitCount+1):
		# 每一层进行FFT的输入数量
		m = 2 ** s
		# 单位W
		wm = Euler(-2 * pi / m)
		# 每层中每个需要进行FFT的初始索引
		for k in range(0, n, m):
			w = complex(1)
			# 对每个FFT进行处理
			for j in range(0, m//2):
				u = rev[k+j]
				t = w * rev[k+j+m//2]
				rev[k+j] = u + t
				rev[k+j+m//2] = u - t
				w = w * wm
	return rev

print('Recursive Version')
for i in RecursiveFFT(xs):
	print(i)
print('Iterative Version')
for i in IterativeFFT(xs):
	print(i)
print('Ground Truth')
for i in fft.fft(xs):
	print(i)

plt.plot(xs, 'r.-')
plt.show()
plt.plot([abs(n) for n in fft.fft(xs)], 'g.-')
plt.show()